import torch
from torchvision import datasets
from torchvision import transforms
from torch import nn
from torch.autograd import Variable
from torch.utils.data import DataLoader

#Lenet5+mnist+底层案例

batch_num = 100
learnning_rate = 0.01

device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')

train = datasets.MNIST(root='MNIST_data',
                       train=True,
                       transform=transforms.ToTensor(),
                       download=True)

test = datasets.MNIST(root='MNIST_data',
                       train=False,
                       transform=transforms.ToTensor(),
                       download=True)

# 管道
dl = DataLoader(dataset=train,
                batch_size=batch_num,
                shuffle=True)

class Lenet5(nn.Module):
    def __init__(self):
        super(Lenet5, self).__init__()
        self.c1m2 = nn.Sequential(
            nn.Conv2d(1, 6, kernel_size=5, stride=1, padding=0), # 24, 24, 6
            nn.ReLU(),
            nn.MaxPool2d(2, 2) # 12, 12, 6
        )
        self.c3m4 = nn.Sequential(
            nn.Conv2d(6, 16, kernel_size=5, stride=1, padding=0), # 8,8,16
            nn.ReLU(),
            nn.MaxPool2d(2, 2) # 4,4,16
        )
        self.fc = nn.Sequential(
            nn.Linear(4*4*16, 120),
            nn.Linear(120, 80),
            nn.Linear(80, 10),
        )
    def forward(self, x):
        out = self.c1m2(x)
        out = self.c3m4(out)
        out = out.view(-1, 4*4*16)
        out = self.fc(out)
        return out

model = Lenet5().to(device)

cost = nn.CrossEntropyLoss()
opti = torch.optim.Adam(model.parameters(), learnning_rate)

for i in range(1):
    avg_cost = 0
    num = len(train) // batch_num
    for j, (x_, y_) in enumerate(dl):
        # 将x，y变为tensor形式
        x = Variable(torch.Tensor(x_)).to(device)
        y = Variable(torch.Tensor(y_.float())).to(device)

        h = model(x)
        c = cost(h, y.long())
        c.backward()
        opti.step()
        opti.zero_grad()
        avg_cost += c/num

    print(i, avg_cost.data.cpu().numpy())

# 测试集准确率
x_test = test.test_data[0:2000].view(2000, 1, 28, 28).float()
y_test = test.test_labels[0:2000]
x_test = Variable(x_test).to(device)
y_test = Variable(y_test).to(device)
h = model(x_test)
pre = torch.max(h, 1)[1]
acc = torch.mean((y_test==pre).float())
print(acc.data.cpu().numpy())